package rpcx_plugins

import (
	"context"
	"net"
	"proxy_resource_test/utils/logging"
	"strings"
	"sync"

	"github.com/smallnest/rpcx/protocol"
	"github.com/smallnest/rpcx/server"
)

type IpMethodCounter struct {
	mu    sync.Mutex
	count map[string]int // key: ip|servicePath|serviceMethod
}

// PostReadRequest 实现 PostReadRequestPlugin 接口,用于统计调用次数
func (p *IpMethodCounter) PostReadRequest(ctx context.Context, req *protocol.Message, err error) error {
	ip := getRemoteIP(ctx)
	if ip == "" {
		ip = "unknown"
	}
	key := ip + "|" + req.ServicePath + "|" + req.ServiceMethod

	p.mu.Lock()
	if p.count == nil {
		p.count = make(map[string]int)
	}
	p.count[key]++
	times := p.count[key]
	p.mu.Unlock()

	logging.Info("[IP统计] %s 调用 %s.%s 第 %d 次", ip, req.ServicePath, req.ServiceMethod, times)
	return nil
}

// getRemoteIP 从 context 获取远端 IP
func getRemoteIP(ctx context.Context) string {
	peer, ok := ctx.Value(server.RemoteConnContextKey).(net.Addr)
	if !ok || peer == nil {
		return ""
	}
	addr := peer.String()
	if i := strings.LastIndex(addr, ":"); i > 0 {
		return addr[:i]
	}
	return addr
}

// 保证实现接口
var _ server.PostReadRequestPlugin = (*IpMethodCounter)(nil)
